{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b2bb7341",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/kumaran/.local/lib/python3.11/site-packages/pydantic/_internal/_config.py:257: UserWarning: Valid config keys have changed in V2:\n",
      "* 'json_encoders' has been removed\n",
      "  warnings.warn(message, UserWarning)\n"
     ]
    }
   ],
   "source": [
    "from datetime import datetime\n",
    "from enum import Enum\n",
    "from os import environ\n",
    "from pathlib import Path\n",
    "from typing import *\n",
    "\n",
    "import yaml\n",
    "from pydantic import BaseModel, Field, NonNegativeInt\n",
    "\n",
    "from fastkafka import FastKafka\n",
    "from fastkafka._components.logger import get_logger"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "132d9f86",
   "metadata": {},
   "outputs": [],
   "source": [
    "logger = get_logger(__name__)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b0303458",
   "metadata": {},
   "outputs": [],
   "source": [
    "class ModelType(str, Enum):\n",
    "    churn = \"churn\"\n",
    "    propensity_to_buy = \"propensity_to_buy\"\n",
    "\n",
    "\n",
    "class ModelTrainingRequest(BaseModel):\n",
    "    AccountId: NonNegativeInt = Field(\n",
    "        ..., example=202020, description=\"ID of an account\"\n",
    "    )\n",
    "    ModelName: ModelType = Field(..., example=\"churn\", description=\"ID of an account\")\n",
    "    total_no_of_records: NonNegativeInt = Field(\n",
    "        ...,\n",
    "        example=1_000_000,\n",
    "        description=\"total number of records (rows) to be ingested\",\n",
    "    )\n",
    "\n",
    "\n",
    "class EventData(BaseModel):\n",
    "    \"\"\"\n",
    "    A sequence of events for a fixed account_id\n",
    "    \"\"\"\n",
    "\n",
    "    AccountId: NonNegativeInt = Field(\n",
    "        ..., example=202020, description=\"ID of an account\"\n",
    "    )\n",
    "    Application: Optional[str] = Field(\n",
    "        None,\n",
    "        example=\"DriverApp\",\n",
    "        description=\"Name of the application in case there is more than one for the AccountId\",\n",
    "    )\n",
    "    DefinitionId: str = Field(\n",
    "        ...,\n",
    "        example=\"appLaunch\",\n",
    "        description=\"name of the event\",\n",
    "        min_length=1,\n",
    "    )\n",
    "    OccurredTime: datetime = Field(\n",
    "        ...,\n",
    "        example=\"2021-03-28T00:34:08\",\n",
    "        description=\"local time of the event\",\n",
    "    )\n",
    "    OccurredTimeTicks: NonNegativeInt = Field(\n",
    "        ...,\n",
    "        example=1616891648496,\n",
    "        description=\"local time of the event as the number of ticks\",\n",
    "    )\n",
    "    PersonId: NonNegativeInt = Field(\n",
    "        ..., example=12345678, description=\"ID of a person\"\n",
    "    )\n",
    "\n",
    "\n",
    "class RealtimeData(BaseModel):\n",
    "    event_data: EventData = Field(\n",
    "        ...,\n",
    "        example=dict(\n",
    "            AccountId=202020,\n",
    "            Application=\"DriverApp\",\n",
    "            DefinitionId=\"appLaunch\",\n",
    "            OccurredTime=\"2021-03-28T00:34:08\",\n",
    "            OccurredTimeTicks=1616891648496,\n",
    "            PersonId=12345678,\n",
    "        ),\n",
    "        description=\"realtime event data\",\n",
    "    )\n",
    "    make_prediction: bool = Field(\n",
    "        ..., example=True, description=\"trigger prediction message in prediction topic\"\n",
    "    )\n",
    "\n",
    "\n",
    "class TrainingDataStatus(BaseModel):\n",
    "    AccountId: NonNegativeInt = Field(\n",
    "        ..., example=202020, description=\"ID of an account\"\n",
    "    )\n",
    "    no_of_records: NonNegativeInt = Field(\n",
    "        ...,\n",
    "        example=12_345,\n",
    "        description=\"number of records (rows) ingested\",\n",
    "    )\n",
    "    total_no_of_records: NonNegativeInt = Field(\n",
    "        ...,\n",
    "        example=1_000_000,\n",
    "        description=\"total number of records (rows) to be ingested\",\n",
    "    )\n",
    "\n",
    "\n",
    "class TrainingModelStatus(BaseModel):\n",
    "    AccountId: NonNegativeInt = Field(\n",
    "        ..., example=202020, description=\"ID of an account\"\n",
    "    )\n",
    "    current_step: NonNegativeInt = Field(\n",
    "        ...,\n",
    "        example=0,\n",
    "        description=\"number of records (rows) ingested\",\n",
    "    )\n",
    "    current_step_percentage: float = Field(\n",
    "        ...,\n",
    "        example=0.21,\n",
    "        description=\"the percentage of the current step completed\",\n",
    "    )\n",
    "    total_no_of_steps: NonNegativeInt = Field(\n",
    "        ...,\n",
    "        example=1_000_000,\n",
    "        description=\"total number of steps for training the model\",\n",
    "    )\n",
    "\n",
    "\n",
    "class ModelMetrics(BaseModel):\n",
    "    \"\"\"The standard metrics for classification models.\n",
    "\n",
    "    The most important metrics is AUC for unbalanced classes such as churn. Metrics such as\n",
    "    accuracy are not very useful since they are easily maximized by outputting the most common\n",
    "    class all the time.\n",
    "    \"\"\"\n",
    "\n",
    "    AccountId: NonNegativeInt = Field(\n",
    "        ..., example=202020, description=\"ID of an account\"\n",
    "    )\n",
    "    Application: Optional[str] = Field(\n",
    "        None,\n",
    "        example=\"DriverApp\",\n",
    "        description=\"Name of the application in case there is more than one for the AccountId\",\n",
    "    )\n",
    "    timestamp: datetime = Field(\n",
    "        ...,\n",
    "        example=\"2021-03-28T00:34:08\",\n",
    "        description=\"UTC time when the model was trained\",\n",
    "    )\n",
    "    mod_type: ModelType = Field(\n",
    "        ...,\n",
    "        example=\"churn\",\n",
    "        description=\"Name of the model used (churn, propensity to buy)\",\n",
    "    )\n",
    "    auc: float = Field(\n",
    "        ..., example=0.91, description=\"Area under ROC curve\", ge=0.0, le=1.0\n",
    "    )\n",
    "    f1: float = Field(..., example=0.89, description=\"F-1 score\", ge=0.0, le=1.0)\n",
    "    precission: float = Field(\n",
    "        ..., example=0.84, description=\"precission\", ge=0.0, le=1.0\n",
    "    )\n",
    "    recall: float = Field(..., example=0.82, description=\"recall\", ge=0.0, le=1.0)\n",
    "    accuracy: float = Field(..., example=0.82, description=\"accuracy\", ge=0.0, le=1.0)\n",
    "\n",
    "\n",
    "class Prediction(BaseModel):\n",
    "    AccountId: NonNegativeInt = Field(\n",
    "        ..., example=202020, description=\"ID of an account\"\n",
    "    )\n",
    "    Application: Optional[str] = Field(\n",
    "        None,\n",
    "        example=\"DriverApp\",\n",
    "        description=\"Name of the application in case there is more than one for the AccountId\",\n",
    "    )\n",
    "    PersonId: NonNegativeInt = Field(\n",
    "        ..., example=12345678, description=\"ID of a person\"\n",
    "    )\n",
    "    prediction_time: datetime = Field(\n",
    "        ...,\n",
    "        example=\"2021-03-28T00:34:08\",\n",
    "        description=\"UTC time of prediction\",\n",
    "    )\n",
    "    mod_type: ModelType = Field(\n",
    "        ...,\n",
    "        example=\"churn\",\n",
    "        description=\"Name of the model used (churn, propensity to buy)\",\n",
    "    )\n",
    "    score: float = Field(\n",
    "        ...,\n",
    "        example=0.4321,\n",
    "        description=\"Prediction score (e.g. the probability of churn in the next 28 days)\",\n",
    "        ge=0.0,\n",
    "        le=1.0,\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "88dc9adc",
   "metadata": {},
   "outputs": [],
   "source": [
    "_total_no_of_records = 0\n",
    "_no_of_records_received = 0\n",
    "\n",
    "\n",
    "def create_ws_server(assets_path: Path = Path(\"./assets\")) -> FastKafka:\n",
    "    title = \"Example for FastKafka\"\n",
    "    description = \"A simple example on how to use FastKafka\"\n",
    "    version = \"0.0.1\"\n",
    "    openapi_url = \"/openapi.json\"\n",
    "    favicon_url = \"/assets/images/favicon.ico\"\n",
    "\n",
    "    contact = dict(name=\"airt.ai\", url=\"https://airt.ai\", email=\"info@airt.ai\")\n",
    "\n",
    "    kafka_brokers = {\n",
    "        \"localhost\": {\n",
    "            \"url\": environ.get(\"KAFKA_HOSTNAME\", \"localhost\"),\n",
    "            \"description\": \"local development kafka\",\n",
    "            \"port\": environ.get(\"KAFKA_PORT\", \"9092\"),\n",
    "        },\n",
    "        \"staging\": {\n",
    "            \"url\": \"kafka.staging.acme.com\",\n",
    "            \"description\": \"staging kafka\",\n",
    "            \"port\": \"9092\",\n",
    "            \"protocol\": \"kafka-secure\",\n",
    "            \"security\": {\"type\": \"plain\"},\n",
    "        },\n",
    "        \"production\": {\n",
    "            \"url\": \"kafka.infobip.acme.com\",\n",
    "            \"description\": \"production kafka\",\n",
    "            \"port\": \"9092\",\n",
    "            \"protocol\": \"kafka-secure\",\n",
    "            \"security\": {\"type\": \"plain\"},\n",
    "        },\n",
    "    }\n",
    "\n",
    "    kafka_server_url = environ.get(\"KAFKA_HOSTNAME\", \"host_not_set\")\n",
    "    kafka_server_port = environ.get(\"KAFKA_PORT\", \"9999\")\n",
    "    kafka_config = {\n",
    "        \"group_id\": f\"{kafka_server_url}:{kafka_server_port}_group\",\n",
    "        \"auto_offset_reset\": \"earliest\",\n",
    "    }\n",
    "    if \"KAFKA_API_KEY\" in environ:\n",
    "        kafka_config = {\n",
    "            **kafka_config,\n",
    "            **{\n",
    "                \"security_protocol\": \"SASL_SSL\",\n",
    "                \"sasl_mechanisms\": \"PLAIN\",\n",
    "                \"sasl_username\": environ[\"KAFKA_API_KEY\"],\n",
    "                \"sasl_password\": environ[\"KAFKA_API_SECRET\"],\n",
    "            },\n",
    "        }\n",
    "\n",
    "    kafka_app = FastKafka(\n",
    "        title=title,\n",
    "        contact=contact,\n",
    "        description=description,\n",
    "        version=version,\n",
    "        kafka_brokers=kafka_brokers,\n",
    "        **kafka_config,\n",
    "    )\n",
    "\n",
    "    @kafka_app.consumes()  # type: ignore\n",
    "    async def on_training_data(msg: EventData):\n",
    "        # ToDo: this is not showing up in logs\n",
    "        logger.debug(f\"msg={msg}\")\n",
    "        global _total_no_of_records\n",
    "        global _no_of_records_received\n",
    "        _no_of_records_received = _no_of_records_received + 1\n",
    "\n",
    "        if _no_of_records_received % 100 == 0:\n",
    "            training_data_status = TrainingDataStatus(\n",
    "                AccountId=EventData.AccountId,\n",
    "                no_of_records=_no_of_records_received,\n",
    "                total_no_of_records=_total_no_of_records,\n",
    "            )\n",
    "            app.produce(\"training_data_status\", training_data_status)\n",
    "\n",
    "    @kafka_app.consumes()  # type: ignore\n",
    "    async def on_realitime_data(msg: RealtimeData):\n",
    "        pass\n",
    "\n",
    "    @kafka_app.produces()  # type: ignore\n",
    "    async def to_training_data_status(msg: TrainingDataStatus) -> TrainingDataStatus:\n",
    "        logger.debug(f\"on_training_data_status(msg={msg}, kafka_msg={kafka_msg})\")\n",
    "        return msg\n",
    "\n",
    "    @kafka_app.produces()  # type: ignore\n",
    "    async def to_training_model_status(msg: str) -> TrainingModelStatus:\n",
    "        logger.debug(f\"on_training_model_status(msg={msg}, kafka_msg={kafka_msg})\")\n",
    "        return TrainingModelStatus()\n",
    "\n",
    "    @kafka_app.produces()  # type: ignore\n",
    "    async def to_model_metrics(msg: ModelMetrics) -> ModelMetrics:\n",
    "        logger.debug(f\"on_training_model_status(msg={msg}, kafka_msg={kafka_msg})\")\n",
    "        return msg\n",
    "\n",
    "    @kafka_app.produces()  # type: ignore\n",
    "    async def to_prediction(msg: Prediction) -> Prediction:\n",
    "        logger.debug(f\"on_realtime_data_status(msg={msg},, kafka_msg={kafka_msg})\")\n",
    "        return msg\n",
    "\n",
    "    return kafka_app"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "03a30de5",
   "metadata": {},
   "outputs": [],
   "source": [
    "def create_app(assets_path: Path = Path(\"../assets\")) -> FastKafka:\n",
    "    assets_path = assets_path.resolve()\n",
    "    kafka_app = create_ws_server(assets_path=assets_path)\n",
    "    return kafka_app"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b49811c7",
   "metadata": {},
   "outputs": [],
   "source": [
    "assets_path: Path = Path(\"../assets\")\n",
    "kafka_app = create_app(assets_path=assets_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8b95a72c",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "python3",
   "language": "python",
   "name": "python3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
